package edu.cn.nlsde.Utils;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;

/**
 * Created by dell on 2016/5/21.
 */
public class SparseMatrix {
    private int row_num;
    private int column_num;
    private HashMap<Integer, HashMap<Integer, Double>> store = new HashMap<>();

    public SparseMatrix() {

    }

    public SparseMatrix(int row, int column) {
        this.row_num = row;
        this.column_num = column;
    }

    public int get_row_num() {
        return row_num;
    }

    public int get_column_num() {
        return column_num;
    }

    public void set(int row, int column, double value) {
        if (!store.containsKey(row)) {
            store.put(row, new HashMap<Integer, Double>());
        }
        store.get(row).put(column, value);
    }

    public double get(int row, int column) {
        HashMap<Integer, Double> tmp = store.get(row);
        if (tmp == null) {
            return 0.0;
        }
        Double value = tmp.get(column);
        if (value == null) {
            return 0.0;
        }
        return value.doubleValue();
    }

    public void add(int row, int column, Integer value) {
        HashMap<Integer, Double> row_list = store.get(row);
        if (row_list == null) {
            row_list = new HashMap<>();
            store.put(row, row_list);
        }
        if (value != 0) {
            Double old_value = row_list.get(column);
            if (old_value == null) {
                row_list.put(column, 1.0 * value);
            } else {
                row_list.put(column, old_value + value);
            }

        }
    }

    public Set<Integer> non_zero() {
        return store.keySet();
    }

    public Set<Integer> non_zero(int key) {
        HashMap<Integer, Double> row_list = store.get(key);
        if (row_list == null) {
            return new HashSet<>();
        }
        return row_list.keySet();
    }

    public void clear() {
        for (int row : store.keySet()) {
            store.get(row).clear();
        }
        store.clear();
    }

}
